package com.yubest.springboothello.config;

import com.yubest.springboothello.interceptors.MdcInterceptor;
import lombok.Getter;
import org.slf4j.MDC;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.core.task.AsyncTaskExecutor;
import org.springframework.scheduling.annotation.EnableAsync;
import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor;

import java.util.UUID;

/**
 * @Author: hweiyu
 * @Date: 2024/8/1 10:02
 */
@Configuration
@EnableAsync
public class AsyncConfig {

    @Bean
    public AsyncTaskExecutor asyncTaskExecutor() {
        ThreadPoolTaskExecutor executor = new ThreadPoolTaskExecutor();
        executor.setCorePoolSize(5);
        executor.setMaxPoolSize(10);
        executor.setQueueCapacity(25);
        executor.setThreadNamePrefix("async-mdc-");
        executor.setTaskDecorator(MdcTaskDecorator::new);
        executor.initialize();
        return executor;
    }

    // 自定义的Runnable装饰器，用于设置MDC
    private static class MdcTaskDecorator implements Runnable {

        @Getter
        private final String tradeId;

        private final Runnable delegate;

        public MdcTaskDecorator(Runnable delegate) {
            this.delegate = delegate;
            String tId = MDC.get(MdcInterceptor.TRACE_ID);
            this.tradeId = null != tId ? tId : UUID.randomUUID().toString();
        }

        @Override
        public void run() {
            MDC.put(MdcInterceptor.TRACE_ID, getTradeId());
            try {
                delegate.run();
            } finally {
                // 清理MDC，避免内存泄漏
                MDC.clear();
            }
        }
    }
}
